"""
人员检测服务
集成YOLOv8模型进行人员检测和跟踪
"""
import asyncio
import logging
import cv2
import numpy as np
from datetime import datetime
from typing import List, Dict, Any, Optional, Tuple
from uuid import uuid4
import time

try:
    from ultralytics import YOLO
    import torch
    YOLO_AVAILABLE = True
except ImportError:
    YOLO_AVAILABLE = False
    logging.warning("YOLOv8 not available. Person detection will use mock data.")

from schemas.person_detection import (
    PersonDetectionResult, PersonBoundingBox, PersonInfo,
    DetectionConfig, DetectionTask, DetectionStatus,
    ModelInfo, PerformanceMetrics
)
from schemas.ai_algorithm import DetectionRequest, DetectionResponse, AIModelConfig, ModelPerformanceMetrics
from enum import Enum
from services.ai_infrastructure.gpu_resource_manager import gpu_resource_manager
from core.config import get_settings

logger = logging.getLogger(__name__)


class ModelStatus(str, Enum):
    """模型状态枚举"""
    LOADING = "loading"
    LOADED = "loaded"
    DISABLED = "disabled"
    ERROR = "error"


class AlgorithmType(str, Enum):
    """算法类型枚举"""
    YOLO = "yolo"
    RCNN = "rcnn"
    SSD = "ssd"
    DETECTRON = "detectron"


class DetectionType(str, Enum):
    """检测类型枚举"""
    PERSON = "person"
    VEHICLE = "vehicle"
    OBJECT = "object"
    FACE = "face"
settings = get_settings()


class PersonTracker:
    """人员跟踪器"""
    
    def __init__(self, max_age: int = 30):
        self.max_age = max_age
        self.tracks: Dict[str, Dict[str, Any]] = {}
        self.next_id = 1
    
    def update(self, detections: List[PersonInfo]) -> List[PersonInfo]:
        """更新跟踪"""
        # 简化的跟踪算法，实际应用中可以使用DeepSORT等更复杂的算法
        updated_detections = []
        
        for detection in detections:
            best_match_id = None
            best_distance = float('inf')
            
            # 查找最佳匹配的跟踪
            for track_id, track in self.tracks.items():
                if track['age'] > self.max_age:
                    continue
                
                # 计算中心点距离
                distance = self._calculate_distance(
                    detection.center, track['last_center']
                )
                
                if distance < best_distance and distance < 100:  # 距离阈值
                    best_distance = distance
                    best_match_id = track_id
            
            if best_match_id:
                # 更新现有跟踪
                detection.tracking_id = best_match_id
                self.tracks[best_match_id]['last_center'] = detection.center
                self.tracks[best_match_id]['age'] = 0
            else:
                # 创建新跟踪
                track_id = f"track_{self.next_id:04d}"
                self.next_id += 1
                detection.tracking_id = track_id
                self.tracks[track_id] = {
                    'last_center': detection.center,
                    'age': 0
                }
            
            updated_detections.append(detection)
        
        # 增加所有跟踪的年龄
        for track in self.tracks.values():
            track['age'] += 1
        
        # 移除过期的跟踪
        expired_tracks = [
            track_id for track_id, track in self.tracks.items()
            if track['age'] > self.max_age
        ]
        for track_id in expired_tracks:
            del self.tracks[track_id]
        
        return updated_detections
    
    def _calculate_distance(self, point1: Tuple[float, float], point2: Tuple[float, float]) -> float:
        """计算两点间距离"""
        return np.sqrt((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2)


class PersonDetectionService:
    """人员检测服务"""
    
    def __init__(self):
        self.models: Dict[str, Any] = {}
        self.trackers: Dict[str, PersonTracker] = {}
        self.detection_tasks: Dict[str, DetectionTask] = {}
        self.performance_metrics: List[PerformanceMetrics] = []
        
        # 默认配置
        self.default_config = DetectionConfig()
        
        # 初始化模型
        self._initialize_models()
    
    def _initialize_models(self):
        """初始化检测模型"""
        if not YOLO_AVAILABLE:
            logger.warning("YOLOv8 not available, using mock model")
            self._initialize_mock_model()
            return
        
        try:
            # 加载YOLOv8模型
            model_configs = [
                {
                    "model_id": "yolov8n_person",
                    "model_path": "yolov8n.pt",
                    "name": "YOLOv8 Nano Person Detection"
                },
                {
                    "model_id": "yolov8s_person", 
                    "model_path": "yolov8s.pt",
                    "name": "YOLOv8 Small Person Detection"
                }
            ]
            
            for config in model_configs:
                try:
                    model = YOLO(config["model_path"])
                    
                    # 注意：不直接修改模型的classes属性，而是在推理时过滤结果
                    
                    self.models[config["model_id"]] = {
                        "model": model,
                        "info": ModelInfo(
                            model_id=config["model_id"],
                            model_name=config["name"],
                            model_type="object_detection",
                            version="8.0.0",
                            framework="ultralytics",
                            input_shape=(3, 640, 640),
                            classes=["person"],
                            model_size=6 if "nano" in config["model_id"] else 14
                        )
                    }
                    
                    logger.info(f"Loaded model: {config['model_id']}")
                    
                except Exception as e:
                    logger.error(f"Failed to load model {config['model_id']}: {e}")
            
            if not self.models:
                logger.warning("No models loaded, using mock model")
                self._initialize_mock_model()
                
        except Exception as e:
            logger.error(f"Failed to initialize models: {e}")
            self._initialize_mock_model()
    
    def _initialize_mock_model(self):
        """初始化模拟模型"""
        self.models["mock_person"] = {
            "model": None,
            "info": ModelInfo(
                model_id="mock_person",
                model_name="Mock Person Detection",
                model_type="object_detection",
                version="1.0.0",
                framework="mock",
                input_shape=(3, 640, 640),
                classes=["person"],
                accuracy=0.85,
                inference_time=20.0,
                model_size=1
            )
        }
    
    async def detect_persons(
        self, 
        image: np.ndarray, 
        camera_id: str,
        config: Optional[DetectionConfig] = None
    ) -> PersonDetectionResult:
        """检测图像中的人员"""
        if config is None:
            config = self.default_config
        
        start_time = time.time()
        detection_id = str(uuid4())
        
        try:
            # 选择模型
            model_id = "yolov8n_person" if "yolov8n_person" in self.models else list(self.models.keys())[0]
            model_data = self.models[model_id]
            
            if model_data["model"] is None:
                # 使用模拟检测
                persons = await self._mock_detection(image, config)
            else:
                # 使用真实模型检测
                persons = await self._real_detection(image, model_data["model"], config)
            
            # 跟踪处理
            if config.enable_tracking:
                if camera_id not in self.trackers:
                    self.trackers[camera_id] = PersonTracker(config.tracking_max_age)
                
                persons = self.trackers[camera_id].update(persons)
            
            # 创建检测结果
            processing_time = time.time() - start_time
            
            result = PersonDetectionResult(
                detection_id=detection_id,
                camera_id=camera_id,
                timestamp=datetime.now(),
                persons=[self._person_to_dict(person) for person in persons],
                total_count=len(persons),
                confidence_threshold=config.confidence_threshold,
                processing_time=processing_time
            )
            
            # 记录性能指标
            await self._record_performance_metrics(model_id, processing_time, len(persons))
            
            return result
            
        except Exception as e:
            logger.error(f"Person detection failed: {e}")
            raise
    
    async def _real_detection(
        self, 
        image: np.ndarray, 
        model: Any, 
        config: DetectionConfig
    ) -> List[PersonInfo]:
        """真实模型检测"""
        try:
            # 预处理图像
            input_image = cv2.resize(image, config.input_size)
            
            # 运行推理
            results = model(input_image, conf=config.confidence_threshold, iou=config.nms_threshold)
            
            persons = []
            for result in results:
                boxes = result.boxes
                if boxes is not None:
                    for i, box in enumerate(boxes):
                        # 检查类别，只处理人员类别（COCO数据集中person类别为0）
                        cls = int(box.cls[0].cpu().numpy())
                        if cls != 0:  # 跳过非人员类别
                            continue
                            
                        # 获取边界框坐标
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                        confidence = float(box.conf[0].cpu().numpy())
                        
                        # 转换回原图尺寸
                        h, w = image.shape[:2]
                        x1 = x1 * w / config.input_size[0]
                        y1 = y1 * h / config.input_size[1]
                        x2 = x2 * w / config.input_size[0]
                        y2 = y2 * h / config.input_size[1]
                        
                        bbox = PersonBoundingBox(
                            x1=x1, y1=y1, x2=x2, y2=y2,
                            confidence=confidence
                        )
                        
                        person = PersonInfo(
                            person_id=f"person_{i}",
                            bbox=bbox,
                            center=bbox.get_center(),
                            features={
                                "width": x2 - x1,
                                "height": y2 - y1,
                                "area": bbox.get_area()
                            }
                        )
                        
                        persons.append(person)
                        
                        if len(persons) >= config.max_detections:
                            break
            
            return persons
            
        except Exception as e:
            logger.error(f"Real detection failed: {e}")
            return []
    
    async def _mock_detection(
        self, 
        image: np.ndarray, 
        config: DetectionConfig
    ) -> List[PersonInfo]:
        """模拟检测"""
        # 模拟延迟
        await asyncio.sleep(0.02)
        
        h, w = image.shape[:2]
        persons = []
        
        # 生成随机数量的人员检测结果
        num_persons = np.random.randint(0, min(10, config.max_detections))
        
        for i in range(num_persons):
            # 随机生成边界框
            x1 = np.random.randint(0, w - 100)
            y1 = np.random.randint(0, h - 200)
            x2 = x1 + np.random.randint(50, 150)
            y2 = y1 + np.random.randint(150, 300)
            
            # 确保边界框在图像范围内
            x2 = min(x2, w)
            y2 = min(y2, h)
            
            confidence = np.random.uniform(config.confidence_threshold, 1.0)
            
            bbox = PersonBoundingBox(
                x1=x1, y1=y1, x2=x2, y2=y2,
                confidence=confidence
            )
            
            person = PersonInfo(
                person_id=f"mock_person_{i}",
                bbox=bbox,
                center=bbox.get_center(),
                features={
                    "width": x2 - x1,
                    "height": y2 - y1,
                    "area": bbox.get_area()
                }
            )
            
            persons.append(person)
        
        return persons
    
    def _person_to_dict(self, person: PersonInfo) -> Dict[str, Any]:
        """将PersonInfo转换为字典"""
        return {
            "person_id": person.person_id,
            "bbox": [person.bbox.x1, person.bbox.y1, person.bbox.x2, person.bbox.y2],
            "confidence": person.bbox.confidence,
            "center": list(person.center),
            "tracking_id": person.tracking_id,
            "features": person.features
        }
    
    async def _record_performance_metrics(
        self, 
        model_id: str, 
        processing_time: float, 
        detection_count: int
    ):
        """记录性能指标"""
        try:
            fps = 1.0 / processing_time if processing_time > 0 else 0
            
            metrics = PerformanceMetrics(
                metric_id=str(uuid4()),
                model_id=model_id,
                timestamp=datetime.now(),
                fps=fps,
                inference_time=processing_time * 1000,  # 转换为毫秒
                preprocessing_time=processing_time * 0.1 * 1000,  # 估算
                postprocessing_time=processing_time * 0.2 * 1000,  # 估算
                memory_usage=512.0,  # 估算值
                gpu_utilization=None  # 可以从GPU资源管理器获取
            )
            
            self.performance_metrics.append(metrics)
            
            # 保持最近1000条记录
            if len(self.performance_metrics) > 1000:
                self.performance_metrics = self.performance_metrics[-1000:]
                
        except Exception as e:
            logger.error(f"Failed to record performance metrics: {e}")
    
    async def create_detection_task(
        self, 
        camera_id: str, 
        task_type: str = "person_detection",
        config: Optional[DetectionConfig] = None
    ) -> DetectionTask:
        """创建检测任务"""
        task_id = str(uuid4())
        
        task = DetectionTask(
            task_id=task_id,
            camera_id=camera_id,
            task_type=task_type,
            status=DetectionStatus.PENDING,
            config=config or self.default_config,
            created_at=datetime.now()
        )
        
        self.detection_tasks[task_id] = task
        
        logger.info(f"Created detection task {task_id} for camera {camera_id}")
        return task
    
    async def start_detection_task(self, task_id: str) -> bool:
        """启动检测任务"""
        if task_id not in self.detection_tasks:
            return False
        
        task = self.detection_tasks[task_id]
        task.status = DetectionStatus.PROCESSING
        task.started_at = datetime.now()
        
        # 这里可以启动实际的检测处理
        # 例如：从视频流中读取帧并进行检测
        
        logger.info(f"Started detection task {task_id}")
        return True
    
    async def get_detection_task(self, task_id: str) -> Optional[DetectionTask]:
        """获取检测任务"""
        return self.detection_tasks.get(task_id)
    
    async def get_detection_tasks(self, camera_id: Optional[str] = None) -> List[DetectionTask]:
        """获取检测任务列表"""
        tasks = list(self.detection_tasks.values())
        
        if camera_id:
            tasks = [task for task in tasks if task.camera_id == camera_id]
        
        return tasks
    
    async def cancel_detection_task(self, task_id: str) -> bool:
        """取消检测任务"""
        if task_id not in self.detection_tasks:
            return False
        
        task = self.detection_tasks[task_id]
        task.status = DetectionStatus.FAILED
        task.completed_at = datetime.now()
        task.error_message = "Task cancelled by user"
        
        logger.info(f"Cancelled detection task {task_id}")
        return True
    
    def get_available_models(self) -> List[ModelInfo]:
        """获取可用模型列表"""
        return [model_data["info"] for model_data in self.models.values()]
    
    def get_model_info(self, model_id: str) -> Optional[ModelInfo]:
        """获取模型信息"""
        if model_id in self.models:
            return self.models[model_id]["info"]
        return None
    
    async def get_performance_metrics(
        self, 
        model_id: Optional[str] = None,
        limit: int = 100
    ) -> List[PerformanceMetrics]:
        """获取性能指标"""
        metrics = self.performance_metrics
        
        if model_id:
            metrics = [m for m in metrics if m.model_id == model_id]
        
        return metrics[-limit:]
    
    async def benchmark_model(self, model_id: str, test_images: List[np.ndarray]) -> Dict[str, Any]:
        """模型基准测试"""
        if model_id not in self.models:
            raise ValueError(f"Model {model_id} not found")
        
        model_data = self.models[model_id]
        config = self.default_config
        
        total_time = 0
        total_detections = 0
        
        for image in test_images:
            start_time = time.time()
            
            if model_data["model"] is None:
                persons = await self._mock_detection(image, config)
            else:
                persons = await self._real_detection(image, model_data["model"], config)
            
            processing_time = time.time() - start_time
            total_time += processing_time
            total_detections += len(persons)
        
        avg_time = total_time / len(test_images)
        avg_fps = 1.0 / avg_time if avg_time > 0 else 0
        avg_detections = total_detections / len(test_images)
        
        return {
            "model_id": model_id,
            "test_images": len(test_images),
            "total_time": total_time,
            "average_time": avg_time,
            "average_fps": avg_fps,
            "average_detections": avg_detections,
            "total_detections": total_detections
        }
    
    def update_detection_config(self, config: DetectionConfig):
        """更新检测配置"""
        self.default_config = config
        logger.info("Updated detection configuration")
    
    def get_detection_config(self) -> DetectionConfig:
        """获取检测配置"""
        return self.default_config
    
    def unload_model(self):
        """卸载模型"""
        if self.model is not None:
            del self.model
            self.model = None
            if self.device and self.device.type == 'cuda':
                torch.cuda.empty_cache()
        
        self.is_loaded = False
        self.config.status = ModelStatus.DISABLED
        logger.info(f"Model {self.config.model_id} unloaded")


# Temporarily commented out duplicate class - needs refactoring
# class PersonDetectionService:
#     """人员检测服务"""
#     
#     def __init__(self):
#         self.models: Dict[str, PersonDetectionModel] = {}
#         self.model_configs: Dict[str, AIModelConfig] = {}
#         self.default_model_id = None
#         self._initialize_default_models()
# 
#     def _initialize_default_models(self):
        """初始化默认模型"""
        # YOLOv8n 人员检测模型
        yolo_config = AIModelConfig(
            model_name="YOLOv8n Person Detection",
            model_type="object_detection",
            version="8.0.0",
            config={
                "model_id": "yolov8n_person",
                "algorithm_type": "YOLO",
                "detection_types": ["PERSON"],
                "model_path": "yolov8n.pt",
                "input_size": [640, 640],
                "confidence_threshold": 0.5,
                "nms_threshold": 0.4,
                "max_detections": 100,
                "gpu_enabled": True,
                "batch_size": 1,
                "created_at": datetime.now().isoformat(),
                "updated_at": datetime.now().isoformat()
            },
            enabled=True
        )
        
        model_id = yolo_config.config["model_id"]
        self.model_configs[model_id] = yolo_config
        self.default_model_id = model_id
    
    async def load_model(self, model_id: str) -> bool:
        """加载模型"""
        try:
            if model_id not in self.model_configs:
                raise ValueError(f"Model config not found: {model_id}")
            
            config = self.model_configs[model_id]
            model = PersonDetectionModel(config)
            
            await model.load_model()
            self.models[model_id] = model
            
            logger.info(f"Model {model_id} loaded successfully")
            return True
            
        except Exception as e:
            logger.error(f"Failed to load model {model_id}: {e}")
            return False
    
    async def unload_model(self, model_id: str) -> bool:
        """卸载模型"""
        try:
            if model_id in self.models:
                self.models[model_id].unload_model()
                del self.models[model_id]
                logger.info(f"Model {model_id} unloaded")
                return True
            return False
        except Exception as e:
            logger.error(f"Failed to unload model {model_id}: {e}")
            return False
    
    async def detect_persons(self, request: DetectionRequest) -> DetectionResponse:
        """检测人员"""
        start_time = time.time()
        request_id = str(uuid4())
        
        try:
            # 选择模型
            model_id = request.model_id
            if model_id not in self.models:
                # 尝试加载模型
                if not await self.load_model(model_id):
                    raise ValueError(f"Failed to load model: {model_id}")
            
            model = self.models[model_id]
            
            # 获取图像
            image = await self._get_image_from_request(request)
            
            # 执行检测
            detections = await model.detect(image)
            
            # 过滤检测类型
            if request.detection_types:
                detections = [d for d in detections if d.detection_type in request.detection_types]
            
            # 过滤置信度
            if request.confidence_threshold:
                detections = [d for d in detections if d.confidence >= request.confidence_threshold]
            
            # 生成标注图像
            annotated_image = None
            if request.return_image:
                annotated_img = model.annotate_image(image, detections)
                # 转换为base64
                _, buffer = cv2.imencode('.jpg', annotated_img)
                annotated_image = base64.b64encode(buffer).decode('utf-8')
            
            processing_time = time.time() - start_time
            
            response = DetectionResponse(
                request_id=request_id,
                camera_id=request.camera_id,
                model_id=model_id,
                detections=detections,
                processing_time=processing_time,
                annotated_image=annotated_image,
                timestamp=datetime.now()
            )
            
            return response
            
        except Exception as e:
            logger.error(f"Person detection failed: {e}")
            raise
    
    async def _get_image_from_request(self, request: DetectionRequest) -> np.ndarray:
        """从请求中获取图像"""
        if request.image_data:
            # 从base64解码
            image_bytes = base64.b64decode(request.image_data)
            nparr = np.frombuffer(image_bytes, np.uint8)
            image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        elif request.image_url:
            # 从URL下载
            response = requests.get(request.image_url)
            response.raise_for_status()
            nparr = np.frombuffer(response.content, np.uint8)
            image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        else:
            raise ValueError("No image data or URL provided")
        
        if image is None:
            raise ValueError("Failed to decode image")
        
        return image
    
    def add_model_config(self, config: AIModelConfig):
        """添加模型配置"""
        model_id = config.config["model_id"]
        self.model_configs[model_id] = config
        logger.info(f"Added model config: {model_id}")
    
    def get_model_config(self, model_id: str) -> Optional[AIModelConfig]:
        """获取模型配置"""
        return self.model_configs.get(model_id)
    
    def get_all_model_configs(self) -> List[AIModelConfig]:
        """获取所有模型配置"""
        return list(self.model_configs.values())
    
    def get_loaded_models(self) -> List[str]:
        """获取已加载的模型列表"""
        return list(self.models.keys())
    
    def get_model_status(self, model_id: str) -> Optional[ModelStatus]:
        """获取模型状态"""
        if model_id in self.models:
            return self.models[model_id].config.status
        elif model_id in self.model_configs:
            return ModelStatus.DISABLED
        return None
    
    def get_model_performance(self, model_id: str) -> Optional[ModelPerformanceMetrics]:
        """获取模型性能指标"""
        if model_id in self.models:
            return self.models[model_id].performance_metrics
        return None
    
    async def benchmark_model(self, model_id: str, test_images: List[np.ndarray]) -> Dict[str, Any]:
        """模型基准测试"""
        if model_id not in self.models:
            if not await self.load_model(model_id):
                raise ValueError(f"Failed to load model: {model_id}")
        
        model = self.models[model_id]
        results = []
        
        for i, image in enumerate(test_images):
            start_time = time.time()
            detections = await model.detect(image)
            processing_time = time.time() - start_time
            
            results.append({
                "image_index": i,
                "processing_time": processing_time,
                "detection_count": len(detections),
                "average_confidence": sum(d.confidence for d in detections) / len(detections) if detections else 0
            })
        
        # 计算统计信息
        total_time = sum(r["processing_time"] for r in results)
        avg_time = total_time / len(results)
        avg_fps = 1.0 / avg_time if avg_time > 0 else 0
        total_detections = sum(r["detection_count"] for r in results)
        avg_detections = total_detections / len(results)
        
        return {
            "model_id": model_id,
            "test_images": len(test_images),
            "total_processing_time": total_time,
            "average_processing_time": avg_time,
            "average_fps": avg_fps,
            "total_detections": total_detections,
            "average_detections": avg_detections,
            "detailed_results": results
        }


# 全局人员检测服务实例
person_detection_service = PersonDetectionService()